#include <Eigen/Dense>
#include <Eigen/Core>
#include <iostream>
#include <fstream>
#include <OsqpEigen/OsqpEigen.h>
#include <string>

#define STATE_NUM   2
#define INPUT_NUM   1
#define PREDICT_WINDOW  10

#define TRACE_FOLLOW    1

void generate_trace(Eigen::Matrix<double, STATE_NUM, 1000> &x_ref_list)
{
    std::ifstream trace_data_file("./trace_data.csv");
    std::vector<double> values;
    std::string line;

    if(!trace_data_file.is_open())
    {
        std::cout << "Error: Open File Failed!" << std::endl;
        std::exit(1);
    }

    int matrix_row_num = 0;

    while (std::getline(trace_data_file, line))
    {
        std::cout << "line: " << std::endl;
        std::stringstream row_line(line);
        std::string num_string;
        std::getline(row_line, num_string, ',');
        int i = 0;
        while(std::getline(row_line, num_string, ','))
        {
            double val = std::stod(num_string);
            std::cout << val << ",";
            values.push_back(val);
            x_ref_list(i, matrix_row_num) = val;
            i++;
        }
        std::cout << std::endl;
        matrix_row_num ++;
    }
    std::cout << x_ref_list << std::endl;

    // x_ref_list = Eigen::Map<const Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>> \
    //             (values.data(), matrix_row_num, values.size()/matrix_row_num);
}

void generate_H_matrix(const Eigen::DiagonalMatrix<double, STATE_NUM> &Q, const Eigen::DiagonalMatrix<double, INPUT_NUM> &R, 
                        int predict_window, Eigen::SparseMatrix<double> &H)
{
    H.resize(STATE_NUM * (predict_window + 1) + INPUT_NUM * predict_window, STATE_NUM * (predict_window + 1) + INPUT_NUM * predict_window);

    for(int i = 0; i < predict_window + 1; i++)
    {
        for(int j = 0; j < STATE_NUM; j++)
        {
            float value = Q.diagonal()[j];
            if(value != 0)
            {
                H.insert(i*STATE_NUM + j, i*STATE_NUM + j) = value;
            }
        }
    }

    for(int i = 0; i < predict_window; i++)
    {
        for(int j = 0; j < INPUT_NUM; j++)
        {
            float value = R.diagonal()[j];
            if(value != 0)
            {
                H.insert((predict_window + 1) * STATE_NUM + i*INPUT_NUM + j, (predict_window + 1) * STATE_NUM + i*INPUT_NUM + j) = value;
            }
        }
    }
}

void generate_f_matrix(const Eigen::DiagonalMatrix<double, STATE_NUM> &Q, const Eigen::Matrix<double, STATE_NUM, 1> &xref, 
                        int predict_window, Eigen::VectorXd &f)
{
    Eigen::Matrix<double, STATE_NUM, 1> Q_xref;
    Q_xref = Q*(-xref);

    f = Eigen::VectorXd::Zero(STATE_NUM * (predict_window + 1) + INPUT_NUM * predict_window, 1);
    for(int i = 0; i < predict_window + 1; i++)
    {
        for(int j = 0; j < STATE_NUM; j++)
        {
            float value = Q_xref(j, 0);
            f(i * STATE_NUM + j, 0) = value;
        }
    }
}

void generate_f_matrix_trace(const Eigen::DiagonalMatrix<double, STATE_NUM> &Q, 
                        const Eigen::Matrix<double, STATE_NUM, PREDICT_WINDOW + 1> &xref, 
                        int predict_window, Eigen::VectorXd &f)
{
    // xref.resize(STATE_NUM, predict_window + 1);
    Eigen::Matrix<double, STATE_NUM, PREDICT_WINDOW + 1> Q_xref;
    Q_xref = Q*(-xref);

    f = Eigen::VectorXd::Zero(STATE_NUM * (predict_window + 1) + INPUT_NUM * predict_window, 1);
    for(int i = 0; i < predict_window + 1; i++)
    {
        for(int j = 0; j < STATE_NUM; j++)
        {
            float value = Q_xref(j, i);
            f(i * STATE_NUM + j, 0) = value;
        }
    }
}

void generate_constraint_matrix(const Eigen::Matrix<double, STATE_NUM, STATE_NUM> &A, const Eigen::Matrix<double, STATE_NUM, INPUT_NUM> &B, 
                    int predict_window, Eigen::SparseMatrix<double> &constraint_matrix)
{
    constraint_matrix.resize(STATE_NUM * (predict_window + 1) * 2 + INPUT_NUM * predict_window, STATE_NUM * (predict_window + 1) + INPUT_NUM * predict_window);
    // -I, equality constriant
    for(int i = 0; i < predict_window + 1; i++)
    {
        for(int j = 0; j < STATE_NUM; j ++)
        {
            constraint_matrix.insert(i*STATE_NUM + j, i*STATE_NUM + j) = -1;
        }
    }

    // A, equality constriant
    for(int i = 0; i < predict_window; i++)
    {
        for(int j = 0; j < STATE_NUM; j++)
        {
            for(int k = 0; k < STATE_NUM; k++)
            {
                if(A(j,k) != 0)
                {
                    constraint_matrix.insert((i+1)*STATE_NUM + j, i*STATE_NUM + k) = A(j,k);
                }
            }                
        }
    }

    // B, equality constriant
    for(int i = 0; i < predict_window; i++)
    {
        for(int j = 0; j < STATE_NUM; j++)
        {
            for(int k = 0; k < INPUT_NUM; k++)
            {
                if(B(j,k) != 0)
                {
                    constraint_matrix.insert((i+1)*STATE_NUM + j, (predict_window + 1)*STATE_NUM + i*INPUT_NUM + k) = B(j,k);
                }
            }                
        }
    }

    // I, inequality constriant
    for(int i = 0; i < STATE_NUM * (predict_window + 1) + INPUT_NUM * predict_window; i++)
    {
        constraint_matrix.insert(STATE_NUM*(predict_window + 1) + i,i) = 1;
    }
}

void generate_constraint_range(const Eigen::Matrix<double, STATE_NUM, 1>&x_max, const Eigen::Matrix<double, STATE_NUM, 1>&x_min, 
                                const Eigen::Matrix<double, INPUT_NUM, 1>&u_max, const Eigen::Matrix<double, INPUT_NUM, 1>&u_min, 
                                const Eigen::Matrix<double, STATE_NUM, 1>&x0, int predict_window,
                                Eigen::VectorXd &low_bound, Eigen::VectorXd &up_bound)
{
    Eigen::VectorXd low_inequality = Eigen::MatrixXd::Zero(STATE_NUM * (predict_window + 1) + INPUT_NUM * (predict_window), 1);
    Eigen::VectorXd up_inequality = Eigen::MatrixXd::Zero(STATE_NUM * (predict_window + 1) + INPUT_NUM * (predict_window), 1);

    for(int i = 0; i < predict_window + 1; i++)
    {
        low_inequality.block(STATE_NUM * i, 0, STATE_NUM, 1) = x_min;
        up_inequality.block(STATE_NUM * i, 0, STATE_NUM, 1) = x_max;
    }
    for(int i = 0; i < predict_window; i++)
    {
        low_inequality.block(STATE_NUM * (predict_window + 1) + i * INPUT_NUM, 0, INPUT_NUM, 1) = u_min;
        up_inequality.block(STATE_NUM * (predict_window + 1) + i * INPUT_NUM, 0, INPUT_NUM, 1) = u_max;
    }

    Eigen::VectorXd low_equality = Eigen::MatrixXd::Zero(STATE_NUM * (predict_window + 1), 1);
    Eigen::VectorXd up_equality = Eigen::MatrixXd::Zero(STATE_NUM * (predict_window + 1), 1);

    low_equality.block(0,0,STATE_NUM,1) = -x0;
    up_equality.block(0,0,STATE_NUM,1) = -x0;

    low_bound = Eigen::MatrixXd::Zero(2 * STATE_NUM * (predict_window + 1) + INPUT_NUM * (predict_window), 1);
    up_bound = Eigen::MatrixXd::Zero(2 * STATE_NUM * (predict_window + 1) + INPUT_NUM * (predict_window), 1);

    low_bound << low_equality, low_inequality;
    up_bound << up_equality, up_inequality;
}

void update_constraint_vector(const Eigen::Matrix<double, STATE_NUM, 1> &x0, Eigen::VectorXd &low_bound, Eigen::VectorXd &up_bound)
{
    low_bound.block(0,0,STATE_NUM, 1) = -x0;
    up_bound.block(0,0,STATE_NUM, 1) = -x0;
}

int main()
{    
    Eigen::Matrix<double, STATE_NUM, STATE_NUM> Ac;
    Eigen::Matrix<double, STATE_NUM, INPUT_NUM> Bc;
    Eigen::Matrix<double, STATE_NUM, STATE_NUM> A;
    Eigen::Matrix<double, STATE_NUM, INPUT_NUM> B;

    int predict_window = PREDICT_WINDOW;

    Eigen::Matrix<double, STATE_NUM, 1> x_max;
    Eigen::Matrix<double, STATE_NUM, 1> x_min;
    Eigen::Matrix<double, INPUT_NUM, 1> u_max;
    Eigen::Matrix<double, INPUT_NUM, 1> u_min;

    Eigen::DiagonalMatrix<double, STATE_NUM> Q;
    Eigen::DiagonalMatrix<double, INPUT_NUM> R;

    Eigen::Matrix<double, STATE_NUM, 1> x0;
    #if TRACE_FOLLOW == 0
    Eigen::Matrix<double, STATE_NUM, 1> x_ref;
    #else
    Eigen::Matrix<double, STATE_NUM, PREDICT_WINDOW+1> predict_x_ref;
    #endif

    Eigen::SparseMatrix<double> H;
    Eigen::VectorXd f;
    Eigen::SparseMatrix<double> constraint_matrix;
    Eigen::VectorXd low_bound;
    Eigen::VectorXd up_bound;

    // 生成轨迹序列
    Eigen::Matrix<double, STATE_NUM, 1000> x_ref_list;
    generate_trace(x_ref_list);

    float m = 1.0f;
    float T = 0.032;

    Ac << 0, 1, 0, 0;
    Bc << 0, 1/m;

    A = Ac * T + Eigen::MatrixXd::Identity(2,2);
    B = Bc * T;

    Q.diagonal() << 1000, 1;
    R.diagonal() << 1;

    x_max << 10, 2;
    x_min << -10, -2;

    u_max << 10;
    u_min << -10;

    x0 << 0, 0;
    #if TRACE_FOLLOW == 0
    x_ref = x_ref_list.block(0, 0, STATE_NUM, 1);
    #else
    predict_x_ref = x_ref_list.block(0, 0, STATE_NUM, predict_window + 1);
    #endif
    
    generate_H_matrix(Q, R, predict_window, H);
    #if TRACE_FOLLOW == 0 
    generate_f_matrix(Q, x_ref, predict_window, f);
    #else
    generate_f_matrix_trace(Q, predict_x_ref, predict_window, f);
    #endif
    generate_constraint_matrix(A, B, predict_window, constraint_matrix);
    generate_constraint_range(x_max, x_min, u_max, u_min, x0, predict_window, low_bound, up_bound);

    OsqpEigen::Solver solver;
    solver.settings()->setWarmStart(true);

    solver.data()->setNumberOfVariables(STATE_NUM*(predict_window + 1) + INPUT_NUM * predict_window);
    solver.data()->setNumberOfConstraints(2 * STATE_NUM * (predict_window + 1) + INPUT_NUM * predict_window);
    if(!solver.data()->setHessianMatrix(H))
    {
        return 1;
    }
    if(!solver.data()->setGradient(f))
    {
        return 1;
    }
    if(!solver.data()->setLinearConstraintsMatrix(constraint_matrix))
    {
        return 1;
    }
    if(!solver.data()->setLowerBound(low_bound))
    {
        return 1;
    }
    if(!solver.data()->setUpperBound(up_bound))
    {
        return 1;
    }
    if(!solver.initSolver())
    {
        return 1;
    }
    
    Eigen::Matrix<double, INPUT_NUM, 1> ctr;
    Eigen::VectorXd QPSolution;

    int num_of_steps = 900;

    std::ofstream out_file;
    out_file.open("./mpc_data.csv");

    for(int i = 0; i < num_of_steps; i++)
    {
        if(solver.solveProblem() != OsqpEigen::ErrorExitFlag::NoError)
        {
            return 1;
        }

        QPSolution = solver.getSolution();

        ctr = QPSolution.block(STATE_NUM * (predict_window + 1), 0, INPUT_NUM, 1);

        auto x0_data = x0.data();
        x0 = A * x0 + B * ctr;
        #if TRACE_FOLLOW == 0
        out_file << i << "," << x0(0) << "," << x0(1) << "," << ctr(0) << "," << x_ref(0,0) << "," << x_ref(1,0) << std::endl;
        #else
        out_file << i << "," << x0(0) << "," << x0(1) << "," << ctr(0) << "," << predict_x_ref(0,0) << "," << predict_x_ref(1,0) << std::endl;
        #endif
        update_constraint_vector(x0, low_bound, up_bound);
        if(!solver.updateBounds(low_bound, up_bound))
        {
            return 1;
        }
        #if TRACE_FOLLOW == 0
        x_ref = x_ref_list.block(0, i, STATE_NUM, 1);        
        generate_f_matrix(Q, x_ref, predict_window, f);
        #else
        predict_x_ref = x_ref_list.block(0, i, STATE_NUM, predict_window + 1);        
        generate_f_matrix_trace(Q, predict_x_ref, predict_window, f);
        #endif
        if(!solver.updateGradient(f))
        {
            return 1;
        }
    }
    
    out_file.close();
    return 0;
}
